import os
import sys
import subprocess
import math
import re
# from openpyxl import Workbook

test_seq_list_8bit=['MarketPlace_1920x1080_60_8bit',
                    'RitualDance_1920x1080_60_8bit',
                    'Cactus_1920x1080_50',
                    'BasketballDrive_1920x1080_50',
                    'BQTerrace_1920x1080_60',
                    'BasketballDrill_832x480_50',
                    'BQMall_832x480_60',
                    'PartyScene_832x480_50',
                    'RaceHorses_832x480_30',
                    'BasketballPass_416x240_50',
                    'BQSquare_416x240_60',
                    'BlowingBubbles_416x240_50',
                    'RaceHorses_416x240_30',
                    'FourPeople_1280x720_60',
                    'Johnny_1280x720_60',
                    'KristenAndSara_1280x720_60']
                    
qp_point=[27,32,38,45]

opt_dict= dict()
# Option: help
opt_dict["-help"]=""
# Option: coding structure
opt_dict["-s"]=""
# Option: YUV path
opt_dict["-i"]=""
# Option: config file path
opt_dict["-c"]=""
# Option: bit depth 10 0:off 1:on default:off
opt_dict["-b"]="0"
# Option: compare rec and dec 0:off 1:on default:on
opt_dict["-m"]="1"
# Option: if run lite test, frame number to be coded is heavily reduced
opt_dict["-l"]="0"
# Option: maximum progresses, defalut fully para
opt_dict["-p"]=str(len(test_seq_list_8bit)*len(qp_point))

def print_help():
    print("Run CTC tests with following options")
    print("\t-s\tcoding structure: AI,RA,LD(LDB)")
    print("\t-i\tYUV path")
    print("\t-c\tconfig file path")
    print("\t-b\tif test 10-bit sequence\t0:off 1:on default:0")
    print("\t-m\tif compare rec and dec\t0:off 1:on default:0")
    print("\t-l\tif run lite test\t0:off 1:on default:0")
    print("\t-p\tmaximum progresses\t[0,%dx%d] default:fully"%(len(test_seq_list_8bit),len(qp_point)))

def get_info_by_name(seq_name:str):
    w=seq_name.split('_')[1].split('x')[0]
    h=seq_name.split('_')[1].split('x')[1]
    f=seq_name.split('_')[2]
    bd=seq_name.split('_')[3].replace('bit','') if len(seq_name.split('_'))==4 else 8
    return int(w),int(h),int(f),int(bd)

def run_tests_in_list(mp:int,test_list):
    batch_lenth=math.ceil(len(test_seq_list_8bit)*len(qp_point)/mp)
    for batch_cnt in range(batch_lenth):
        process_num=mp if not batch_cnt==batch_lenth-1 else len(test_seq_list_8bit)*len(qp_point)- mp*(batch_lenth-1)
        p_list=list()
        for process_cnt in range(process_num):
            print("[START] "+str(test_list[batch_cnt*mp + process_cnt][0])+"")
            p_temp=subprocess.Popen(str(test_list[batch_cnt*mp + process_cnt][0]),shell=True)
            p_list.append(p_temp)
        for process in p_list:
            #p_temp:subprocess.Popen = process
            p_temp = process
            p_temp.wait()
        for process_cnt in range(process_num):
            print("[START] "+str(test_list[batch_cnt*mp + process_cnt][1])+"")
            print("[START] "+str(test_list[batch_cnt*mp + process_cnt][2])+"")
            p_temp=subprocess.Popen(str(test_list[batch_cnt*mp + process_cnt][1]),shell=True)
            p_temp.wait()
            p_temp=subprocess.Popen(str(test_list[batch_cnt*mp + process_cnt][2]),shell=True)
            p_temp.wait()


def get_result(log_list):
    print("[START] collect results.")
    attr_list=["bitrate","PSNR Y","PSNR U","PSNR V"]
    f=open("./result.txt",'w')
    for log in log_list:
        has_result=0
        with open(log, 'r',encoding='utf8') as file_to_read:
            while True:
                lines = file_to_read.readline() # 整行读取数
                if not lines:
                    if has_result==0:
                        f.write('no result for ' + log + '\n')
                    break
                elif "PSNR Y" in lines:
                    y_psnr = float(re.findall(r"\d+\.\d+|\d+", lines)[0])
                    has_result=1
                elif "PSNR U" in lines:
                    u_psnr = float(re.findall(r"\d+\.\d+|\d+", lines)[0])
                    has_result=1
                elif "PSNR V" in  lines:
                    v_psnr = float(re.findall(r"\d+\.\d+|\d+", lines)[0])
                    has_result=1
                elif "bitrate(kbps)" in  lines:
                    bitrate = float(re.findall(r"\d+\.\d+|\d+", lines)[0])  
                    has_result=1
                elif "Total encoding time" in lines:
                    enc_time = float(re.findall(r"\d+\.\d+|\d+", lines)[1])  
                    has_result=1

        if  has_result==1:
            results = (log  + '\t'    
                        + '\t' + (str('%.4f' % y_psnr))
                        + '\t' + (str('%.4f' % u_psnr))  
                        + '\t' + (str('%.4f' % v_psnr)) 
                        + '\t' + (str('%.4f' % bitrate)) 
                        + '\t' + (str('%.4f' % enc_time)) 
                        + '\n')
            f.write(results)

    f.close()

def test_ctc(cs:str,yuv_path:str,config_path:str,test_10_bit:bool,is_lite_test:bool,compare_rec_dec:bool,mp:int):
    
    test_list=list()
    log_list=list()
    lite_fn=10
    if cs=='ra':
        frame_num_list=[120,100,60,120,100,100,120,120,120,120,120,120,120,100,60,120,100,100,120,120,120,120,120,120,120,120]
        basic_enc_line="../build_cmake/app/app_encoder -c %s/encoder_randomaccess_vtm.cfg -i %s -wdt %d -hgt %d -fr %d -ip %d -f %d -q %d -b bin/test_%s.bin -o yuv/test_%s_rec.yuv > log/test_%s.txt"
        itra_ra=[64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64]
    elif cs=='ldb' or cs=='ld':
        frame_num_list=[40,40,40,40,40,100,100,100,100,120,120,120,120,100,100,100,100,100,120,120,120,120,120,120,120,120]
        basic_enc_line="../build_cmake/app/app_encoder --config %s/encode_IPPP.cfg -i %s -w %d -h %d -z %d -f %d -d 8 -q %d -o bin/test_%s.bin -r yuv/test_%s_rec.yuv > log/test_%s.txt"
        itra_ra=[64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64,64]
    elif cs=='ai':
        frame_num_list=[120,100,60,120,100,100,120,120,120,120,120,120,120,100,60,120,100,100,120,120,120,120,120,120,120,120]
        basic_enc_line="../bin/app_encoder --config %s/encode_AI.cfg -i %s -w %d -h %d -z %d -f %d -d 8 -q %d -o bin/test_%s.bin -r yuv/test_%s_rec.yuv > log/test_%s.txt"
        itra_ra=[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]
        lite_fn=160
    else:
        print("Unknown Coding structure %s !"%cs)
        return 
    
    basic_dec_line="../bin/app_decoder -s -i bin/test_%s.bin -o yuv/test_%s_dec.yuv"

    for seq_idx in range(len(test_seq_list_8bit)):
        seq_path=os.path.join(yuv_path,str(test_seq_list_8bit[seq_idx]))+".yuv"
        if test_10_bit:
            seq_path=seq_path.replace('8bit','10bit')
        if not os.path.exists(seq_path):
            print(seq_path+"do not exist !\n")
            return 
        ip=itra_ra[seq_idx]
        f_num=frame_num_list[seq_idx] if not is_lite_test else min(lite_fn,frame_num_list[seq_idx])
        width,height,fr,bit_depth=get_info_by_name(str(test_seq_list_8bit[seq_idx]))
        for qp in qp_point:
            test_name=str(test_seq_list_8bit[seq_idx])+'_'+str(qp)
            enc = basic_enc_line%(config_path,seq_path,width,height,fr,f_num,qp,test_name,test_name,test_name)
            print(basic_enc_line)
            dec = basic_dec_line%(test_name,test_name)
            cmp =("./cmp_yuv.sh yuv/test_%s_rec.yuv yuv/test_%s_dec.yuv"%(test_name,test_name)) if compare_rec_dec else ("rm yuv/test_%s_rec.yuv yuv/test_%s_dec.yuv"%(test_name,test_name))
            test_list.append((enc,dec,cmp))
            log_list.append("log/test_%s.txt"%test_name)

    run_tests_in_list(mp,test_list)
    get_result(log_list)
    print("[FINISHED]!!")

if __name__ == "__main__":

    invalid_or_help=0
    if len(sys.argv)>1:
        lastOpt=""
        for i in range(1,len(sys.argv)):
            if str(sys.argv[i]) in opt_dict:
                lastOpt=str(sys.argv[i])
                if lastOpt=="-help":
                    invalid_or_help=1
                    break
            elif len(lastOpt)>0:
                opt_dict[lastOpt]=str(sys.argv[i])
                lastOpt=""
        if len(opt_dict["-s"])>0 and len(opt_dict["-i"])>0 and len(opt_dict["-c"])>0:
            coding_structure=opt_dict["-s"].lower()
            print()
        else:
            invalid_or_help=1
    else:
        invalid_or_help=1
    
    if(invalid_or_help):
        print_help()
    else:
        test_ctc(coding_structure,str(opt_dict['-i']),str(opt_dict['-c']),int(opt_dict['-b'])==1,int(opt_dict['-l'])==1,int(opt_dict['-m'])==1,int(opt_dict['-p']))
